package com.zzyy.study.config;

import com.zzyy.study.service.ChatAssistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.jina.JinaScoringModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.router.DefaultQueryRouter;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * @auther zzyy
 * @create 2025-03-16 22:31
 */

@Configuration(proxyBeanMethods = false)
public class ChatLanguageModelConfig
{

    @Bean
    public ChatLanguageModel chatLanguageModel()
    {
        return OpenAiChatModel.builder()
                .apiKey(System.getenv("LANGCHAIN4J_KEY"))
                .modelName("qwen-max")
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .build();
    }

    /**
     * 需要预处理文档并将其存储在专门的嵌入存储（也称为矢量数据库）中。当用户提出问题时，这对于快速找到相关信息是必要的。
     * 我们可以使用我们支持的 15 多个嵌入存储中的任何一个，但为了简单起见，我们将使用内存中的嵌入存储：
     * @return
     */
    @Bean
    public InMemoryEmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }


    @Bean
    public ChatAssistant assistant(ChatLanguageModel chatLanguageModel, EmbeddingStore<TextSegment> embeddingStore)
    {

        DefaultRetrievalAugmentor defaultRetrievalAugmentor = DefaultRetrievalAugmentor.builder()
                //查询转换器(queryTransformer)配合使用一个LLM压缩查询和对话上下文，生成更好的查询
                .queryTransformer(new CompressingQueryTransformer(chatLanguageModel))
                //内容检索器(contentRetriever)代表根据用户的查询从底层数据源中获取内容
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                //查询路由器(QueryRouter)负责将查询分配到合适的内容检索器，可以并存多个内容源
                .queryRouter(new DefaultQueryRouter(EmbeddingStoreContentRetriever.from(embeddingStore)))
                // 注入内容
                .contentInjector(new DefaultContentInjector())
                //并行化处理，当只有一个查询和一个检索器时，DefaultRetrievalAugmentor
                // 会在单线程中执行查询和检索。否则，它会使用 Executor 并行处理多个查询和检索任务
                .executor(new ThreadPoolExecutor(Runtime.getRuntime().availableProcessors(),
                        Runtime.getRuntime().availableProcessors() * 2,
                        1L,
                        TimeUnit.SECONDS,
                        new LinkedBlockingQueue<>(1000)))
                .build();



        return AiServices.builder(ChatAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .retrievalAugmentor(defaultRetrievalAugmentor)
                .build();
    }
}
